# -*- coding: utf-8 -*-
# time: 2025/3/29 10:55
# file: ch01.py
# author: hanson
"""
Embedding Models：chromadb 内置的 all-MiniLM-L6-v2

"""
import chromadb.config
import ollama

# 方式一：直接创建 client 数据都存储在内存，程序运行完，数据会丢失 测试使用
client = chromadb.Client()
# 文档定义
documents = [
    "咖啡树是一种茜草科的常绿灌木或小乔木，原产于非洲的埃塞俄比亚和苏丹",
    "咖啡树的果实被称为咖啡樱桃，成熟时会变成鲜红色或紫色，每个果实通常包含两颗咖啡豆",
    "咖啡豆是咖啡树果实内的种子，经过烘焙后可以研磨成粉末冲泡咖啡",
    "阿拉比卡咖啡和罗布斯塔咖啡是两种主要的咖啡品种，阿拉比卡咖啡口感更佳，而罗布斯塔咖啡咖啡因含量更高",
    "咖啡种植需要温暖的气候和充足的降雨，咖啡树通常生长在海拔1000米到2000米的山区",
    "咖啡是世界上最受欢迎的饮料之一，具有提神醒脑的作用，并且富含抗氧化剂",
]
collection = client.create_collection(name="docs")
# 将每个文档存储在向量嵌入数据库中,在添加documents的时候，我们也可以补充metadatas和embeddings / ids 为后续提供更多可用数据。
for i,w in enumerate(documents):
    collection.add(documents=[w], metadatas=[{"source": i}],ids=[str(i)]);

#检索prompt的相关信息
prompt = "咖啡树的果实被称为什么？"
results = collection.query(
    query_texts=[prompt],
    n_results=3, # 检索的文档数量
)
print(results)
print("-------------------------")
# out:{'ids': [['1', '2', '5']], 'distances': [[0.7086873054504395, 0.7665070295333862, 1.0334540605545044]], 'metadatas':
#其中distances是近似度，我们可以根据这个来判断相关性，在 ChromaDB 中，distances 值越小，代表相似度越高。
"""
现实使用过程中相似性查询不是特别准确，我们一般会和全文索引结果结合起来一起评估。
把检索出来的相关性数据作为prompt的一部分提供给AI进行对话，输出数据，我这里用的是qwen2.5:7b，
当然，这里不一定使用ollama，你可以使用其它任意大模型API
本地的Qwen2___5-1___5B-Instruct模型导入Ollama, ollama 用户需要创建一个Modelfile文件
ollama create your-model-name -f your-path/Modelfile
"""

data = results["documents"][0][0]
print(data)
print("-------------------------")
# 这里不一定使用ollama，你可以使用其它任意大模型API，替换以下代码即可
# deepseek-r1:1.5b
output = ollama.generate(
    model="deepseek-r1:1.5b",
    #model="qwen2.5:1.5b",
    prompt=f"根据这段文字：{data}。回答这个问题：{prompt}",
)
print("-------------------------")
print(output["response"])



